# src/plot_data/comparecase.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*

import os
import json
import argparse
import sys 
import pandas as pd
import numpy as np

from baseplot import plot_linechart
from baseplot import plot_percent_2D
from baseplot import plot_percent
from baseplot import plot_percent_line
from baseplot import plot_heatmap
from baseplot import casename2number
from baseplot import getplteffdata


sys.path.append("..")
from package.plot_utility import get_runcases

# set parser
def make_parser(): 
    parser = argparse.ArgumentParser(
        prog="jxpamg-predo-merge-pid-timerdata",
        description="jxpamg timer data preprocess script, merge pid's timer data to a single json.")
    parser.add_argument(
            "-sdir","--savedir",
            dest="savedir",
            default="./",
            help="Destination save pic/csv directories."
        )
    parser.add_argument(
            "-po", "--picopt",
            dest="plot_type",
            default="time_linechart",
            help="Set type of pic out. Option: rank, time_heatmap, time_linechart, percent_par, percent_total, percent_line, efficiency, sizeperprocess."
    )
    parser.add_argument(
            "-co", "--csvopt",
            dest="csv_type",
            default="time",
            help="Set type of data in csv. Option: time, rank_correlation."
    )
    parser.add_argument(
            "-o", "--output",
            dest="out_type",
            default="csv",
            help="Set output type of plotting picture. Option: csv, pic."
    )
    parser.add_argument(
            "LOGDIR",
            metavar="LOGDIR",
            help="Destination log directories."
        )
    parser.add_argument(
            "COMPONENT",
            metavar="COMPONENT",
            default="CreateS",
            help="Set component of plotting picture. Option: ."
    )
    return parser

# dataprocessing
# 获得不同 case 得到 component 对应的 time
def getjsondata(casename, component, plttype):
    timedic = {}
    for i in casename:
        basename = os.path.basename(i)
        basename = basename.split("-")[0]
        with open(i, "r") as fo:
            text = fo.read()
            jsondata = json.loads(text)
            print(f"Load {basename} data.")
        timedic[basename] = getavetime(jsondata, component, plttype)

    return timedic

# 获取对应 level 和 parent 的 模块 的 ave 时间
def getavetime(jsondata, component, plttype):
# TODO:
# ave time 获取 和 total time 分开
    datalist = []
    for i in jsondata:
        # 加入np
        tempdic = {}
        if i == "dir":
            continue
        tempdic["np"] = int(i.split("_")[1].strip("np"))
        if(plttype == "percent_line"):
            for j in jsondata[i]:
                if j == component :
                    tempdic[j] = float(jsondata[i][j]["tavg"])
                    root = jsondata[i][j]["root"]
            for j in jsondata[i]:
                if j == root :
                    tempdic["TotalSolveTime"] = float(jsondata[i][j]["tavg"])
        else:
            #paren == NULL,取得level相同的组件
            for j in jsondata[i]:
                if j == component :
                    tempdic[j] = float(jsondata[i][j]["tavg"])
                elif j == "Totaltime" :
                    tempdic["TotalSolveTime"] = float(jsondata[i][j]["tavg"])
        datalist.append(tempdic)
    # print(datalist)
    return datalist

# 导出为CSV
def savedata2csv(savedir, casename, component, csv_type, datadict):
    caselist = casename2number(casename)
    writer = pd.ExcelWriter(savedir + f"Compare_{component}_{csv_type}.xlsx")

    savedatadict = {}

    k = 0
    for i in datadict:
        data = pd.DataFrame(datadict[i])
        if csv_type == "time":
            data = data.sort_index(axis=1)
            data = data.sort_values(by='np', ascending=True)
        savedatadict[f"c{caselist[k]}-{component}"] = data[component]
        k+=1
    if csv_type == "time":
        savedatadict["np"] = data["np"]
    savedatadict = pd.DataFrame(savedatadict)
    savedatadict.to_excel(writer,sheet_name=component)
    writer.save()

# csv 数据处理

def save2csv(savedir,runcases,component,csv_type,datadict):
    """表格数据处理"""
    if csv_type == "time":
        savedata2csv(savedir, runcases, component, csv_type, datadict)
    elif csv_type == "rank_correlation":
        """
        参考文章Statistical Scalability Analysis of Communication Operations in Distributed Applications的指标
        1. 先把不同进程数下的rank算出来
        2. 再把rank与时间做线性相关
        """
        resultdic = {}
        for i in datadict:
            temp_ditc = {}
            data = pd.DataFrame(datadict[i])
            data = data.sort_index(axis=1)
            data = data.sort_values(by='np', ascending=True)
            data = data.drop(labels="TotalSolveTime",axis=1)
            data = data.copy(deep=True)
            for col in list(data.columns):
                data[col] = data[col].rank()
            for col in list(data.columns):
                if(col != "np"):
                    
                    rs = np.cov(data["np"].values[:],data[col].values[:])[0][1] / (np.std(data["np"].values[:]) * np.std(data[col].values[:]))
                    temp_ditc[col] = rs
            resultdic[i] = [temp_ditc]
        savedata2csv(savedir,runcases, component, csv_type, resultdic)
    return

# 画图处理的代码

# 得到画图的输入数据
# TODO:这个代码没有存在的必要，下面这些数据都可以通过pandas获得
def getpltdata(data) :
    x = data['np'].values.tolist()
    legend_name = []
    for i in data :
        if i != "np" :
            legend_name.append(i)
    y = []
    for i in legend_name :
        y.append(list(data[i]))
    return (x,y,legend_name)

# plot
def save2pic(savedir,casename,component,datadict,plttype):
    savedatadict = {}
    Totaltimedict = {}
    casenumber = casename2number(casename)

    for i in casename:
        basename = os.path.basename(i)
        basename = basename.split("-")[0]
        data = pd.DataFrame(datadict[basename])
        data = data.sort_index(axis=1)
        data = data.sort_values(by='np', ascending=True)

        # if(plttype == "percent_line"):
        #     for i in list(data.columns):
        #         if i == "Setup" or i == "Solve":
        #             Totaltimedict[f"{basename}"] = data[i]
        # else:
        #     Totaltimedict[f"{basename}"] = data["TotalSolveTime"]
        Totaltimedict[f"{basename}"] = data["TotalSolveTime"]
        savedatadict[f"{basename}"] = data[component]

    #TODO:下面的代码写的太繁琐了，需要修改
    savedatadict["np"] = data["np"]
    casenumber.append("np")
    savedatadict = pd.DataFrame(savedatadict)
    Totaltimedict = pd.DataFrame(Totaltimedict)
    
    # 删除部分行代码
    # savedatadict = savedatadict.drop(labels=[7],axis=0)
    # Totaltimedict = Totaltimedict.drop(labels=[7],axis=0)
    
    savedatadict.columns = casenumber
    casenumber.remove("np")
    Totaltimedict.columns = casenumber
    x,y,legend_name = getpltdata(savedatadict)
    title = f"{component}"
    if plttype == "time_heatmap" :
        # print(savedatadict.columns)
        # print(savedatadict.columns)
        savedatadict.set_index(["np"], inplace=True)
        plot_heatmap(savedir, component + "_timeheatmap", "casenumber", savedatadict.T)
        # data = data.drop(labels="TotalSolveTime",axis=1)
        # data.set_index(["np"], inplace=True)
        # plot_heatmap(savedir,title + "-heatmap",plttype,data.T)
    elif plttype == "rank_heatmap" :
        savedatadict.set_index(["np"], inplace=True)
        savedatadict = savedatadict.T
        for i in list(savedatadict.columns):
            savedatadict[i] = savedatadict[i].rank()
        plot_heatmap(savedir, component + "_rankheatmap", "casenumber", savedatadict)
    elif plttype == "time_linechart" :
        plot_linechart(savedir,x,y,f"{plttype}(s)",legend_name,title)
    elif plttype == "percent_par" :
        plot_percent(savedir,x,y,plttype,legend_name,title)
    elif plttype == "percent_total" :
        plot_percent_2D(savedir,x,savedatadict,plttype,legend_name,title,Totaltimedict)
    elif plttype == "percent_line" :
        plot_percent_line(savedir,x,savedatadict,plttype,legend_name,title,Totaltimedict)
    elif plttype == "efficiency" :
        x,y = getplteffdata(x,y)
        plot_linechart(savedir,x,y,f"{plttype}(%)",legend_name,title+"-efficiency")
    return

def main():
    parser = make_parser()
    args = parser.parse_args()
    runcases = get_runcases(args.LOGDIR)
    datadict = getjsondata(runcases, args.COMPONENT, args.plot_type)
    # print(pd.DataFrame(datadict))
    # for i in datadict:
    #     print(pd.DataFrame(datadict[i]))

    if(args.out_type == "csv"):
        save2csv(args.savedir, runcases, args.COMPONENT, args.csv_type, datadict)
    if(args.out_type == "pic"):
        save2pic(args.savedir, runcases, args.COMPONENT, datadict,args.plot_type)

if __name__ == "__main__":
    main()
    # test()
